import h5py
import os

def print_hdf5_structure(file_path, group_path="", show_details=False):
    """
    打印HDF5文件的结构信息
    
    参数:
    file_path: HDF5文件路径
    group_path: 要查看的组路径，默认为空(根目录)
    show_details: 是否显示详细信息
    """
    
    def print_item_info(name, obj, indent_level=0):
        """打印单个项的信息"""
        indent = "  " * indent_level
        
        if isinstance(obj, h5py.Dataset):
            print(f"{indent}DATASET: {name}")
            print(f"{indent}  Shape: {obj.shape}")
            print(f"{indent}  Data type: {obj.dtype}")
            if show_details and obj.size > 0 and obj.size <= 100:
                try:
                    print(f"{indent}  Data sample: {obj[()]}")
                except:
                    print(f"{indent}  Data sample: <无法显示>")
        elif isinstance(obj, h5py.Group):
            # 统计子项
            n_datasets = 0
            n_groups = 0
            for item_name in obj.keys():
                if isinstance(obj[item_name], h5py.Dataset):
                    n_datasets += 1
                elif isinstance(obj[item_name], h5py.Group):
                    n_groups += 1
            
            print(f"{indent}GROUP: {name} (子组: {n_groups}, 数据集: {n_datasets})")
    
    try:
        with h5py.File(file_path, 'r') as f:
            if group_path == "":
                # 打印顶层结构
                print(f"HDF5文件: {os.path.basename(file_path)}")
                print("=" * 50)
                
                n_datasets = 0
                n_groups = 0
                
                # 遍历顶层项
                for key in f.keys():
                    obj = f[key]
                    print_item_info(key, obj, 0)
                    
                    # 统计顶层项数量
                    if isinstance(obj, h5py.Dataset):
                        n_datasets += 1
                    elif isinstance(obj, h5py.Group):
                        n_groups += 1
                
                print(f"\n总计: {n_groups} 个组, {n_datasets} 个数据集")
            else:
                # 打印指定组下的结构
                if group_path in f:
                    group = f[group_path]
                    if isinstance(group, h5py.Group):
                        print(f"HDF5文件: {os.path.basename(file_path)} - 组: {group_path}")
                        print("=" * 50)
                        
                        n_datasets = 0
                        n_groups = 0
                        
                        # 遍历指定组下的项
                        for key in group.keys():
                            obj = group[key]
                            print_item_info(key, obj, 0)
                            
                            # 统计项数量
                            if isinstance(obj, h5py.Dataset):
                                n_datasets += 1
                            elif isinstance(obj, h5py.Group):
                                n_groups += 1
                        
                        print(f"\n总计: {n_groups} 个组, {n_datasets} 个数据集")
                    else:
                        print(f"路径 '{group_path}' 不是一个组")
                else:
                    print(f"组 '{group_path}' 不存在")
                    
    except FileNotFoundError:
        print(f"错误: 文件 '{file_path}' 不存在")
    except Exception as e:
        print(f"错误: {e}")

def explore_hdf5(file_path, path=""):
    """
    交互式探索HDF5文件结构
    
    参数:
    file_path: HDF5文件路径
    path: 初始查看路径
    """
    print_hdf5_structure(file_path, path)
    
    with h5py.File(file_path, 'r') as f:
        current_path = path
        while True:
            print(f"\n当前路径: {'/' if current_path == '' else current_path}")
            command = input("输入组名进入子组，'..' 返回上级，'quit' 退出: ").strip()
            
            if command == 'quit':
                break
            elif command == '..':
                # 返回上级目录
                if current_path != "":
                    parts = current_path.split('/')
                    current_path = '/'.join(parts[:-1]) if len(parts) > 1 else ""
            elif command in f:
                # 进入子组
                obj = f[command] if current_path == "" else f[f"{current_path}/{command}"]
                if isinstance(obj, h5py.Group):
                    current_path = command if current_path == "" else f"{current_path}/{command}"
                else:
                    print(f"'{command}' 不是一个组")
            elif command in [key for key in (f[current_path].keys() if current_path else f.keys())]:
                # 如果在当前路径下找到了该组
                obj = f[command] if current_path == "" else f[f"{current_path}/{command}"]
                if isinstance(obj, h5py.Group):
                    current_path = command if current_path == "" else f"{current_path}/{command}"
                else:
                    print(f"'{command}' 不是一个组")
            else:
                print(f"未找到 '{command}'")
            
            # 清屏并显示当前路径下的内容
            print("\n" + "-"*50)
            print_hdf5_structure(file_path, current_path)

# 使用示例
if __name__ == "__main__":
    # 替换为你的HDF5文件路径
    hdf5_file_path = "D:\projects\datasets\hwdb.h5"
    
    # # 1. 打印顶层结构
    # print_hdf5_structure(hdf5_file_path)
    
    # print("\n" + "="*60 + "\n")
    
    # # 2. 打印特定组下的结构
    # print_hdf5_structure(hdf5_file_path, "train")
    
    # print("\n" + "="*60 + "\n")

    # # 3. 显示详细信息
    # print_hdf5_structure(hdf5_file_path, "train/Gnt1.0TrainPart1")
    
    # print("\n" + "="*60 + "\n")

    # # 3. 显示详细信息
    # print_hdf5_structure(hdf5_file_path, "", show_details=True)
    
    # 4. 交互式探索（取消注释以使用）
    explore_hdf5(hdf5_file_path)